#!/usr/bin/env python3
"""
XGBoost模型调优器 - 针对电力负荷预测优化
解决当前XGBoost模型R²=0.4068的欠拟合问题
"""

import pandas as pd
import numpy as np
from datetime import datetime, date, timedelta
from typing import Dict, List, Optional, Tuple, Any
import logging
import os
import pickle
import json
from pathlib import Path

# 机器学习模型
import xgboost as xgb
from sklearn.multioutput import MultiOutputRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split, GridSearchCV, cross_val_score
from sklearn.pipeline import Pipeline

# 数据库相关
from sqlalchemy.orm import Session
from sqlalchemy import text, and_

from backend.config.database import get_db_session
from backend.entities.load_data_new import LoadData
from backend.entities.weather_daily import WeatherDaily
from backend.entities.holiday import HolidayInfo
from backend.entities.model_training_log import ModelTrainingLog
from backend.entities.model_eval import ModelEval
from backend.utils.holiday_rule_engine import is_holiday

logger = logging.getLogger(__name__)

class XGBoostTuner:
    """XGBoost模型调优器"""
    
    def __init__(self):
        self.model_dir = Path("backend/models/saved_models")
        self.model_dir.mkdir(parents=True, exist_ok=True)
        
        # 当前性能差的配置（R² = 0.4068）
        self.current_config = {
            'objective': 'reg:squarederror',
            'eval_metric': 'rmse',
            'max_depth': 4,  # 太浅
            'learning_rate': 0.05,  # 太低
            'n_estimators': 150,  # 太少
            'subsample': 0.6,  # 太低
            'colsample_bytree': 0.6,  # 太低
            'colsample_bylevel': 0.6,
            'colsample_bynode': 0.6,
            'gamma': 5.0,  # 太高，过度限制分裂
            'min_child_weight': 3,
            'reg_alpha': 0.1,  # L1正则化太强
            'reg_lambda': 0.1,  # L2正则化太强
            'random_state': 42,
            'n_jobs': -1,
            'verbosity': 0
        }
        
        # 基础优化配置 - 提高模型复杂度
        self.optimized_config = {
            'objective': 'reg:squarederror',
            'eval_metric': 'rmse',
            'max_depth': 8,  # 增加深度提高拟合能力
            'learning_rate': 0.1,  # 提高学习率
            'n_estimators': 500,  # 大幅增加树数量
            'subsample': 0.85,  # 提高采样率
            'colsample_bytree': 0.8,  # 提高特征采样
            'colsample_bylevel': 0.8,
            'colsample_bynode': 0.8,
            'gamma': 0.1,  # 降低分裂阈值
            'min_child_weight': 1,  # 降低叶子节点权重限制
            'reg_alpha': 0.01,  # 降低L1正则化
            'reg_lambda': 0.01,  # 降低L2正则化
            'random_state': 42,
            'n_jobs': -1,
            'verbosity': 0,
            'tree_method': 'hist',  # 使用histogram方法
            'grow_policy': 'depthwise'  # 深度优先增长
        }
        
        # 高级配置候选 - 用于网格搜索
        self.param_grids = {
            'aggressive': {
                'max_depth': [6, 8, 10, 12],
                'learning_rate': [0.08, 0.1, 0.15, 0.2],
                'n_estimators': [300, 500, 800, 1000],
                'subsample': [0.8, 0.85, 0.9],
                'colsample_bytree': [0.8, 0.85, 0.9],
                'gamma': [0, 0.1, 0.2],
                'min_child_weight': [1, 2, 3],
                'reg_alpha': [0, 0.01, 0.05],
                'reg_lambda': [0, 0.01, 0.05]
            },
            'conservative': {
                'max_depth': [5, 6, 7],
                'learning_rate': [0.05, 0.08, 0.1],
                'n_estimators': [200, 300, 500],
                'subsample': [0.7, 0.8, 0.85],
                'colsample_bytree': [0.7, 0.8, 0.85],
                'gamma': [0.1, 0.2, 0.5],
                'min_child_weight': [1, 2],
                'reg_alpha': [0.01, 0.05],
                'reg_lambda': [0.01, 0.05]
            }
        }
        
        # 特征工程策略
        self.feature_strategies = {
            'basic': self._basic_feature_engineering,
            'advanced': self._advanced_feature_engineering,
            'ensemble': self._ensemble_feature_engineering
        }
    
    def load_data(self, train_range: str = "3_months") -> Tuple[np.ndarray, np.ndarray]:
        """加载训练数据"""
        try:
            # 确定数据范围
            if train_range == "1_month":
                days_back = 30
            elif train_range == "2_months":
                days_back = 60
            elif train_range == "3_months":
                days_back = 90
            elif train_range == "6_months":
                days_back = 180
            else:
                days_back = 90
            
            end_date = datetime.now().date()
            start_date = end_date - timedelta(days=days_back)
            
            logger.info(f"加载数据范围: {start_date} 到 {end_date}")
            
            with get_db_session() as db:
                # 获取负荷数据
                load_query = db.query(LoadData).filter(
                    LoadData.dt >= start_date,
                    LoadData.dt <= end_date
                ).order_by(LoadData.dt, LoadData.t_idx)
                
                load_data = load_query.all()
                
                if not load_data:
                    raise ValueError("没有找到负荷数据")
                
                # 按日期组织数据
                daily_loads = {}
                for record in load_data:
                    dt_key = record.dt
                    if dt_key not in daily_loads:
                        daily_loads[dt_key] = [0.0] * 96
                    if 0 <= record.t_idx <= 95:
                        daily_loads[dt_key][record.t_idx] = record.load_val
                
                # 获取完整天数的数据
                complete_days = []
                for dt, loads in daily_loads.items():
                    if all(load > 0 for load in loads):  # 确保数据完整
                        complete_days.append((dt, loads))
                
                complete_days.sort(key=lambda x: x[0])
                logger.info(f"找到 {len(complete_days)} 天完整数据")
                
                if len(complete_days) < 14:  # 至少需要14天数据
                    raise ValueError(f"完整数据不足，仅有 {len(complete_days)} 天")
                
                # 准备特征和标签
                X_list = []
                y_list = []
                
                for i in range(7, len(complete_days)):  # 需要前7天作为特征
                    target_date, target_loads = complete_days[i]
                    
                    # 历史负荷特征（前7天）
                    historical_loads = []
                    for j in range(7):
                        hist_date, hist_loads = complete_days[i-7+j]
                        historical_loads.extend(hist_loads)
                    
                    # 气象特征
                    weather_data = db.query(WeatherDaily).filter(WeatherDaily.dt == target_date).first()
                    if weather_data:
                        weather_features = [
                            weather_data.t_max or 20.0,
                            weather_data.t_min or 10.0,
                            weather_data.precip or 0.0,
                            weather_data.humidity or 60.0
                        ]
                    else:
                        weather_features = [20.0, 10.0, 0.0, 60.0]
                    
                    # 时间特征
                    weekday = target_date.weekday()
                    month = target_date.month
                    is_holiday_flag = is_holiday(target_date, db)
                    
                    time_features = [
                        1.0 if weekday == i else 0.0 for i in range(7)
                    ] + [
                        month / 12.0,
                        1.0 if is_holiday_flag else 0.0
                    ]
                    
                    # 组合特征
                    feature_vector = historical_loads + weather_features + time_features
                    
                    X_list.append(feature_vector)
                    y_list.append(target_loads)
                
                X = np.array(X_list)
                y = np.array(y_list)
                
                logger.info(f"数据加载完成: X.shape={X.shape}, y.shape={y.shape}")
                return X, y
                
        except Exception as e:
            logger.error(f"数据加载失败: {e}")
            raise
    
    def _basic_feature_engineering(self, X: np.ndarray) -> np.ndarray:
        """基础特征工程"""
        return X  # 使用原始特征
    
    def _advanced_feature_engineering(self, X: np.ndarray) -> np.ndarray:
        """高级特征工程"""
        try:
            X_engineered = X.copy()
            
            # 历史负荷统计特征
            if X.shape[1] >= 672:  # 确保有历史负荷数据
                historical_loads = X[:, :672]  # 前672列是历史负荷
                
                # 添加统计特征
                load_mean = np.mean(historical_loads, axis=1, keepdims=True)
                load_std = np.std(historical_loads, axis=1, keepdims=True)
                load_max = np.max(historical_loads, axis=1, keepdims=True)
                load_min = np.min(historical_loads, axis=1, keepdims=True)
                load_median = np.median(historical_loads, axis=1, keepdims=True)
                
                # 趋势特征
                load_trend = []
                for i in range(X.shape[0]):
                    recent_loads = historical_loads[i, -96:]  # 最近24小时
                    time_indices = np.arange(96)
                    trend = np.corrcoef(time_indices, recent_loads)[0, 1]
                    load_trend.append(trend if not np.isnan(trend) else 0.0)
                load_trend = np.array(load_trend).reshape(-1, 1)
                
                # 周期性特征
                daily_patterns = []
                for i in range(X.shape[0]):
                    # 计算每天同一时间点的平均值
                    daily_avg = []
                    for t in range(96):
                        same_time_loads = historical_loads[i, t::96]  # 7天中同一时间点
                        daily_avg.append(np.mean(same_time_loads))
                    daily_patterns.extend(daily_avg)
                daily_patterns = np.array(daily_patterns).reshape(X.shape[0], -1)
                
                X_engineered = np.column_stack([
                    X_engineered, 
                    load_mean, 
                    load_std, 
                    load_max, 
                    load_min,
                    load_median,
                    load_trend,
                    daily_patterns
                ])
            
            logger.info(f"高级特征工程：{X.shape[1]} -> {X_engineered.shape[1]} 维")
            return X_engineered
            
        except Exception as e:
            logger.warning(f"高级特征工程失败: {e}")
            return X
    
    def _ensemble_feature_engineering(self, X: np.ndarray) -> np.ndarray:
        """集成特征工程 - 结合多种策略"""
        try:
            # 先应用高级特征工程
            X_advanced = self._advanced_feature_engineering(X)
            
            # 添加交互特征
            if X.shape[1] >= 676:  # 有气象特征
                weather_features = X[:, 672:676]  # 气象特征
                load_features = X[:, :672]  # 历史负荷
                
                # 气象-负荷交互特征
                temp_max = weather_features[:, 0].reshape(-1, 1)
                temp_min = weather_features[:, 1].reshape(-1, 1)
                humidity = weather_features[:, 3].reshape(-1, 1)
                
                recent_load_avg = np.mean(load_features[:, -96:], axis=1, keepdims=True)
                
                # 温度与负荷的交互
                temp_load_interaction = temp_max * recent_load_avg / 10000  # 归一化
                humidity_load_interaction = humidity * recent_load_avg / 100000
                
                X_advanced = np.column_stack([
                    X_advanced,
                    temp_load_interaction,
                    humidity_load_interaction
                ])
            
            logger.info(f"集成特征工程：{X.shape[1]} -> {X_advanced.shape[1]} 维")
            return X_advanced
            
        except Exception as e:
            logger.warning(f"集成特征工程失败: {e}")
            return self._advanced_feature_engineering(X)
    
    def evaluate_config(self, config: Dict[str, Any], X: np.ndarray, y: np.ndarray, 
                       feature_strategy: str = 'basic') -> Dict[str, float]:
        """评估特定配置的性能"""
        try:
            # 应用特征工程
            X_engineered = self.feature_strategies[feature_strategy](X)
            
            # 划分训练测试集
            X_train, X_test, y_train, y_test = train_test_split(
                X_engineered, y, test_size=0.2, random_state=42
            )
            
            # 创建模型
            base_model = xgb.XGBRegressor(**config)
            model = MultiOutputRegressor(base_model, n_jobs=-1)
            
            # 训练模型
            model.fit(X_train, y_train)
            
            # 预测
            train_pred = model.predict(X_train)
            test_pred = model.predict(X_test)
            
            # 计算指标
            train_mae = mean_absolute_error(y_train, train_pred)
            test_mae = mean_absolute_error(y_test, test_pred)
            train_rmse = np.sqrt(mean_squared_error(y_train, train_pred))
            test_rmse = np.sqrt(mean_squared_error(y_test, test_pred))
            train_r2 = r2_score(y_train, train_pred)
            test_r2 = r2_score(y_test, test_pred)
            
            # 过拟合评分
            overfitting_score = (train_r2 - test_r2) / max(train_r2, 0.001)
            
            metrics = {
                'train_mae': train_mae,
                'test_mae': test_mae,
                'train_rmse': train_rmse,
                'test_rmse': test_rmse,
                'train_r2': train_r2,
                'test_r2': test_r2,
                'overfitting_score': overfitting_score,
                'feature_strategy': feature_strategy,
                'feature_count': X_engineered.shape[1]
            }
            
            logger.info(f"配置评估完成 - R²: {test_r2:.4f}, MAE: {test_mae:.2f}")
            return metrics
            
        except Exception as e:
            logger.error(f"配置评估失败: {e}")
            return {
                'test_r2': 0.0,
                'test_mae': 9999.0,
                'overfitting_score': 1.0
            }
    
    def tune_xgboost(self, train_range: str = "3_months") -> Dict[str, Any]:
        """XGBoost调优主函数"""
        try:
            logger.info("开始XGBoost调优...")
            
            # 加载数据
            X, y = self.load_data(train_range)
            
            # 评估当前配置
            logger.info("评估当前配置...")
            current_metrics = self.evaluate_config(self.current_config, X, y, 'basic')
            logger.info(f"当前配置性能: R²={current_metrics['test_r2']:.4f}, MAE={current_metrics['test_mae']:.2f}")
            
            # 评估基础优化配置
            logger.info("评估基础优化配置...")
            optimized_metrics = self.evaluate_config(self.optimized_config, X, y, 'basic')
            logger.info(f"优化配置性能: R²={optimized_metrics['test_r2']:.4f}, MAE={optimized_metrics['test_mae']:.2f}")
            
            # 测试不同特征工程策略
            best_config = self.optimized_config.copy()
            best_metrics = optimized_metrics
            best_strategy = 'basic'
            
            for strategy in ['advanced', 'ensemble']:
                logger.info(f"测试特征工程策略: {strategy}")
                strategy_metrics = self.evaluate_config(self.optimized_config, X, y, strategy)
                logger.info(f"{strategy}策略性能: R²={strategy_metrics['test_r2']:.4f}, MAE={strategy_metrics['test_mae']:.2f}")
                
                if strategy_metrics['test_r2'] > best_metrics['test_r2']:
                    best_metrics = strategy_metrics
                    best_strategy = strategy
                    logger.info(f"找到更好的策略: {strategy}")
            
            # 进行网格搜索优化
            logger.info("开始网格搜索优化...")
            grid_results = self._grid_search_optimization(X, y, best_strategy)
            
            if grid_results['test_r2'] > best_metrics['test_r2']:
                best_config = grid_results['best_config']
                best_metrics = grid_results
                logger.info(f"网格搜索找到更好配置: R²={best_metrics['test_r2']:.4f}")
            
            # 保存最佳配置
            self._save_best_config(best_config, best_metrics, best_strategy)
            
            logger.info("XGBoost调优完成!")
            logger.info(f"最佳性能: R²={best_metrics['test_r2']:.4f}, MAE={best_metrics['test_mae']:.2f}")
            logger.info(f"性能提升: R²提升 {best_metrics['test_r2'] - current_metrics['test_r2']:.4f}")
            
            return {
                'success': True,
                'current_metrics': current_metrics,
                'best_metrics': best_metrics,
                'best_config': best_config,
                'best_strategy': best_strategy,
                'improvement': {
                    'r2_improvement': best_metrics['test_r2'] - current_metrics['test_r2'],
                    'mae_improvement': current_metrics['test_mae'] - best_metrics['test_mae']
                }
            }
            
        except Exception as e:
            logger.error(f"XGBoost调优失败: {e}")
            return {'success': False, 'error': str(e)}
    
    def _grid_search_optimization(self, X: np.ndarray, y: np.ndarray, 
                                 feature_strategy: str) -> Dict[str, Any]:
        """网格搜索优化"""
        try:
            # 应用特征工程
            X_engineered = self.feature_strategies[feature_strategy](X)
            
            # 准备网格搜索参数（简化版，避免计算量过大）
            param_grid = {
                'max_depth': [6, 8, 10],
                'learning_rate': [0.08, 0.1, 0.15],
                'n_estimators': [300, 500],
                'subsample': [0.8, 0.85],
                'colsample_bytree': [0.8, 0.85]
            }
            
            best_score = -np.inf
            best_params = None
            
            logger.info(f"网格搜索参数空间大小: {np.prod([len(v) for v in param_grid.values()])}")
            
            # 简化的网格搜索
            for max_depth in param_grid['max_depth']:
                for learning_rate in param_grid['learning_rate']:
                    for n_estimators in param_grid['n_estimators']:
                        config = self.optimized_config.copy()
                        config.update({
                            'max_depth': max_depth,
                            'learning_rate': learning_rate,
                            'n_estimators': n_estimators
                        })
                        
                        metrics = self.evaluate_config(config, X, y, feature_strategy)
                        score = metrics['test_r2']
                        
                        if score > best_score:
                            best_score = score
                            best_params = config
                            logger.info(f"新的最佳配置: R²={score:.4f}, depth={max_depth}, lr={learning_rate}, n_est={n_estimators}")
            
            if best_params:
                # 用最佳参数重新评估
                best_metrics = self.evaluate_config(best_params, X, y, feature_strategy)
                best_metrics['best_config'] = best_params
                return best_metrics
            else:
                return {'test_r2': 0.0, 'best_config': self.optimized_config}
                
        except Exception as e:
            logger.error(f"网格搜索失败: {e}")
            return {'test_r2': 0.0, 'best_config': self.optimized_config}
    
    def _save_best_config(self, config: Dict[str, Any], metrics: Dict[str, Any], 
                         strategy: str):
        """保存最佳配置"""
        try:
            # 保存配置到文件
            config_path = self.model_dir / 'xgboost_tuned_config.json'
            with open(config_path, 'w') as f:
                json.dump({
                    'config': config,
                    'metrics': {k: float(v) if isinstance(v, (np.float64, np.float32)) else v 
                              for k, v in metrics.items()},
                    'feature_strategy': strategy,
                    'tuning_time': datetime.now().isoformat()
                }, f, indent=2)
            
            logger.info(f"最佳配置已保存到: {config_path}")
            
        except Exception as e:
            logger.error(f"保存配置失败: {e}")


# 使用示例
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    tuner = XGBoostTuner()
    results = tuner.tune_xgboost("3_months")
    
    if results['success']:
        print(f"调优成功!")
        print(f"R²提升: {results['improvement']['r2_improvement']:.4f}")
        print(f"MAE改善: {results['improvement']['mae_improvement']:.2f}")
    else:
        print(f"调优失败: {results['error']}")